'''
Image Classification Training Database
By BrokenData
'''

from PIL import Image
import numpy as np

class TrainingGroups():
    def __init__(self,Resolution) -> None:
        '''
        Resolution : Resolution of the image.
        '''
        self.ImageList=[]
        self.Generated=False
        self.Resolution=Resolution
    def AddImage(self,image_pointer):
        #Checking the type of the image.
        if len(self.Resolution) == 3:
            COLOR=True
        elif len(self.Resolution) == 2:
            COLOR=False
        else:
            raise Exception('It seemed that the resolution is wrong! The resolution is %s'%self.Resolution)
        #So let's change the image into a grey mode.
        if not COLOR:
            im=image_pointer.convert('L')
        else:
            im=image_pointer.copy()
        #Resize the image
        im=im.resize((self.Resolution[0],self.Resolution[1]))
        self.ImageList.append(np.asarray(im))
    def AddImageFile(self,image_path):
        ImagePointer=Image.open(image_path)
        self.AddImage(ImagePointer)
    def GiveImageGroup(self):
        return self.ImageList


class ImageTrainingDatabase():
    def __init__(self,Kinds) -> None:
        '''
        Kinds : kinds of the image
        '''
        if type(Kinds) != list:
            raise Exception('We hope the type of the image is a list,but %s is given.'%type(Kinds))
        if len(set(Kinds)) != len(Kinds):
            raise Exception('We can\'t accept a list with items having same names!')
        self.Kinds=list(Kinds)
        self.KindsAdded=[False]*len(Kinds)
        self.TrainImage=[]
        self.TrainKinds=[]
        self.KindLen=len(self.Kinds)
    def AddGroup(self,ImageGroup,TrainKindName):
        KindNum=self.Kinds.index(TrainKindName)
        for Im in ImageGroup:
            self.TrainImage.append(Im)
            KindList=[0]*self.KindLen
            KindList[KindNum]=1
            self.TrainKinds.append(KindList)
    def GiveTrainData(self):
        return (np.asarray(self.TrainImage),np.asarray(self.TrainKinds))
